/**
 * Copyright 2024 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "plugin/device/ascend/kernel/opapi/aclnn/argmax_ext_aclnn_kernel.h"
#include <algorithm>
#include <vector>
#include <memory>
#include <functional>
#include "ir/tensor.h"
#include "runtime/device/kernel_runtime.h"
#include "transform/acl_ir/op_api_convert.h"
#include "abstract/ops/primitive_infer_map.h"

namespace mindspore {
namespace kernel {

void ArgMaxAscend::GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs,
                                    const std::vector<KernelTensor *> &outputs) {
  dim_ = 0;
  keepdim_ = false;
  auto in_shape = inputs[kIndex0]->GetShapeVector();
  input_realshape_ = in_shape;
  output_realshape_ = outputs[kIndex0]->GetShapeVector();
  auto dim_value_opt = inputs[kIndex1]->GetOptionalValueWithCheck<int64_t>();
  if (dim_value_opt.has_value()) {
    dim_ = dim_value_opt.value();
    keepdim_ = transform::ConvertKernelTensor<bool>(inputs[kIndex2]);
  } else {  // input dim is None set flatten size
    int input_flatten_size = std::accumulate(in_shape.begin(), in_shape.end(), 1, std::multiplies<int64_t>());
    input_realshape_ = ShapeVector{input_flatten_size};
    inputs[kIndex0]->SetShapeVector(input_realshape_);
  }
  GetWorkspaceForResize(inputs[kIndex0], dim_, keepdim_, outputs[kIndex0]);
}

bool ArgMaxAscend::Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
                          const std::vector<KernelTensor *> &outputs, void *stream_ptr) {
  MS_EXCEPTION_IF_NULL(stream_ptr);
  ParseGenExecutor(GEN_EXECUTOR_BOOST(op_type_, hash_id_, inputs[kIndex0], dim_, keepdim_, outputs[kIndex0]));
  RunOp(stream_ptr, workspace);
  return true;
}

MS_ACLNN_KERNEL_FACTORY_REG(ArgMaxExt, ArgMaxAscend);
}  // namespace kernel
}  // namespace mindspore
